#include <linux/init.h>
#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/completion.h>
#include <linux/kthread.h>
#include <linux/fs.h>

#include <linux/module.h>
#include <linux/irq.h>
#include <linux/platform_device.h>
#include <linux/dma-mapping.h>
#include <linux/cdev.h>
#include <linux/types.h>
#include <linux/mutex.h>
#include <linux/acpi.h>
#include <crypto/internal/hash.h>
#include <linux/crypto.h>
#include <crypto/hash.h>
#include <crypto/internal/skcipher.h>
#include <crypto/internal/akcipher.h>
#include <linux/ioport.h>
#include <linux/io.h>

#include "sm4_test.h"
#include "test.h"

static int sm4_enc_test(const u8 *key, u8 *iv, const u8 *data, unsigned int len,
			u8 *out)
{
	struct crypto_wait wait;
	struct skcipher_request *req;
	struct crypto_skcipher *tfm;
	struct scatterlist ssg, dsg;
	static u32 flag = 0;
	int ret = -1;
	u32 offset, step_len;
	u8 *tmp = NULL;

	tfm = crypto_alloc_skcipher(alg, 0, 0);
	if (IS_ERR(tfm)) {
		pr_err("failed to load transform for %s: %ld\n", alg,
		       PTR_ERR(tfm));
		return ret;
	}

	if (flag++ == 0)
		pr_info("\ntesting of async %s (%s) enc!\n", alg,
			get_driver_name(crypto_skcipher, tfm));

	tmp = kmalloc(len, GFP_KERNEL);
	if (tmp == NULL)
		goto out;
	memcpy(tmp, data, len);

	req = skcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		pr_info("skcipher: Failed to allocate request for %s\n",
			"ctr(sm4)");
		goto out;
	}

	crypto_skcipher_clear_flags(tfm, ~0);
	crypto_skcipher_setkey(tfm, key, 16);

	for (offset = 0; offset < len;) {
		if (step) {
			step_len = 32;
			if (offset + step_len > len)
				step_len = len - offset;

		} else {
			step_len = len;
		}

		crypto_init_wait(&wait);
		skcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					      crypto_req_done, &wait);

		sg_init_table(&ssg, 1);
		sg_set_buf(&ssg, tmp + offset, step_len);
		sg_init_table(&dsg, 1);
		sg_set_buf(&dsg, out + offset, step_len);

		skcipher_request_set_crypt(req, &ssg, &dsg, step_len, iv);

		ret = crypto_wait_req(crypto_skcipher_encrypt(req), &wait);
		if (ret) {
			pr_err("crypto_skcipher_encrypt failed ret=%d\n", ret);
			break;
		}
		offset += step_len;
	}

out:
	if (tmp)
		kfree(tmp);

	if (req)
		skcipher_request_free(req);

	if (tfm)
		crypto_free_skcipher(tfm);

	return ret;
}

static int sm4_dec_test(const u8 *key, u8 *iv, const u8 *data, unsigned int len,
			u8 *out)
{
	struct crypto_wait wait;
	struct skcipher_request *req;
	struct crypto_skcipher *tfm;
	struct scatterlist ssg, dsg;
	static u32 flag = 0;
	int ret;
	u32 offset, step_len;
	u8 *tmp = NULL;

	tfm = crypto_alloc_skcipher(alg, 0, 0);
	if (IS_ERR(tfm)) {
		pr_err("failed to load transform for %s: %ld\n", alg,
		       PTR_ERR(tfm));
		return ret;
	}

	if (flag++ == 0)
		pr_info("\ntesting of async %s (%s) dec!\n", alg,
			get_driver_name(crypto_skcipher, tfm));

	tmp = kmalloc(len, GFP_KERNEL);
	if (tmp == NULL)
		goto out;

	memcpy(tmp, data, len);

	req = skcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		pr_info("skcipher: Failed to allocate request for %s\n", alg);
		goto out;
	}

	crypto_skcipher_clear_flags(tfm, ~0);
	if (key != NULL)
		crypto_skcipher_setkey(tfm, key, 16);

	for (offset = 0; offset < len;) {
		if (step) {
			step_len = 32;
			if (offset + step_len > len)
				step_len = len - offset;
		} else {
			step_len = len;
		}

		crypto_init_wait(&wait);
		skcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					      crypto_req_done, &wait);

		sg_init_table(&ssg, 1);
		sg_set_buf(&ssg, tmp + offset, step_len);
		sg_init_table(&dsg, 1);
		sg_set_buf(&dsg, out + offset, step_len);

		skcipher_request_set_crypt(req, &ssg, &dsg, step_len, iv);

		ret = crypto_wait_req(crypto_skcipher_decrypt(req), &wait);
		if (ret) {
			pr_err("crypto_skcipher_encrypt failed ret=%d\n", ret);
			break;
		}
		offset += step_len;
	}

out:
	if (tmp)
		kfree(tmp);

	if (req)
		skcipher_request_free(req);

	if (tfm)
		crypto_free_skcipher(tfm);

	return ret;
}

/* 0:ecb 1: cbc 2: ctr */
void sm4_test(int mode)
{
	u8 *result, *src;
	u32 maxlen = 2048;
	int i, count, ret;
	struct cipher_testvec *testvecs = NULL;
	u8 tmpiv[16];
	u8 *iv_ptr = NULL;

	if (mode == 0) {
		count = sizeof(sm4_tv_template) / sizeof(sm4_tv_template[0]);
		testvecs = (struct cipher_testvec *)sm4_tv_template;
	} else if (mode == 1) {
		count = sizeof(sm4_cbc_tv_template) /
			sizeof(sm4_cbc_tv_template[0]);
		testvecs = (struct cipher_testvec *)sm4_cbc_tv_template;
	} else if (mode == 2) {
		count = sizeof(sm4_ctr_tv_template) /
			sizeof(sm4_ctr_tv_template[0]);
		testvecs = (struct cipher_testvec *)sm4_ctr_tv_template;
	} else {
		pr_err("Unsupport mode\n");
		return;
	}

	src = kmalloc(maxlen, GFP_KERNEL);
	if (src == NULL)
		return;

	result = kmalloc(maxlen, GFP_KERNEL);
	if (result == NULL) {
		kfree(src);
		return;
	}

	pr_info("get vector count %d\n", count);
	for (i = 0; i < count; i++) {
		if (idx != 0xffff && i != idx) {
			continue;
		}

		if (testvecs[i].len > maxlen) {
			pr_err("data len %d > maxlen %d\n", testvecs[i].len,
			       maxlen);
			break;
		}

		if (testvecs[i].iv) {
			memcpy(tmpiv, testvecs[i].iv, 16);
			iv_ptr = tmpiv;
		} else
			iv_ptr = NULL;

		/* test encrypt */
		memcpy(src, testvecs[i].ptext, testvecs[i].len);
		ret = sm4_enc_test(testvecs[i].key, iv_ptr, src,
				   testvecs[i].len, result);
		if (ret) {
			break;
		}

		// compare it
		if (memcmp(testvecs[i].ctext, result, testvecs[i].len)) {
			pr_err("enc %s err!\n", alg);
			print_buf("key", (u8 *)testvecs[i].key, 16);
			if (iv_ptr) {
				print_buf("iv", (u8 *)testvecs[i].iv, 16);
				print_buf("outiv", iv_ptr, 16);
			}
			print_buf("src", src, testvecs[i].len);
			print_buf("result", result, testvecs[i].len);
			print_buf("expected ctext", (u8 *)testvecs[i].ctext,
				  testvecs[i].len);

			break;
		}
		pr_info("Test alg %s [%d] enc success\r\n", alg, i);

		if (testvecs[i].iv) {
			memcpy(tmpiv, testvecs[i].iv, 16);
			iv_ptr = tmpiv;
		} else
			iv_ptr = NULL;

		memset(src, 0, testvecs[i].len);
		ret = sm4_dec_test(testvecs[i].key, iv_ptr, result,
				   testvecs[i].len, src);
		if (ret) {
			break;
		}

		// compare it
		if (memcmp(testvecs[i].ptext, src, testvecs[i].len)) {
			pr_err("dec %s err!\n", alg);
			print_buf("key", (u8 *)testvecs[i].key, 16);
			if (iv_ptr) {
				print_buf("iv", (u8 *)testvecs[i].iv, 16);
				print_buf("outiv", iv_ptr, 16);
			}
			print_buf("ctext", result, testvecs[i].len);
			print_buf("result", src, testvecs[i].len);
			print_buf("expected ptext", (u8 *)testvecs[i].ptext,
				  testvecs[i].len);

			break;
		}

		pr_info("Test alg %s [%d] dec success\r\n", alg, i);
	}

	kfree(src);
	kfree(result);
}
